package cn.mrdjun.redis.stater.ratelimiter;

import cn.mrdjun.redis.stater.util.AspectUtil;
import cn.mrdjun.redis.stater.util.IpUtil;
import com.google.common.collect.ImmutableList;
import lombok.RequiredArgsConstructor;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.Objects;

/**
 * 单IP限流 - 防止恶意调用
 *
 * @author MrDJun
 */
@Aspect
@RequiredArgsConstructor
public class IpRateLimitAspect {
    private final static Logger log = LoggerFactory.getLogger(IpRateLimitAspect.class);
    private final RedisTemplate<String, Object> redisTemplate;

    @Pointcut("@annotation(cn.mrdjun.redis.stater.ratelimiter.IpRateLimit)")
    public void pointcut() {
    }

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        HttpServletRequest request = getHttpServletRequest();
        Method method = AspectUtil.resolveMethod(point);
        IpRateLimit limitAnnotation = method.getAnnotation(IpRateLimit.class);

        String name = limitAnnotation.name();
        int limitPeriod = limitAnnotation.period();
        int limitCount = limitAnnotation.count();

        String ip = IpUtil.getIpAddr(request);
        String uri = request.getRequestURI();
        ImmutableList<String> keys = ImmutableList.of(uri + "_" + ip);
        RedisScript<Long> redisScript = new DefaultRedisScript<>(TEXT, Long.class);
        Long count = redisTemplate.execute(redisScript, keys, limitCount, limitPeriod);
        if (count != null && count.intValue() <= limitCount) {
            log.debug("IP[{}]第[{}]次访问接口[{}]", ip, count, uri);
            return point.proceed();
        } else {
            log.warn("超出频率限制：IP[{}]接口[{}]描述[{}]", ip, uri, name);
            return false;
        }
    }

    /**
     * 限流脚本
     * 调用的时候不超过阈值，则直接返回并执行计算器自加。
     */
    private final static String TEXT = "local c" +
            "\nc = redis.call('get',KEYS[1])" +
            "\nif c and tonumber(c) > tonumber(ARGV[1]) then" +
            "\nreturn c;" +
            "\nend" +
            "\nc = redis.call('incr',KEYS[1])" +
            "\nif tonumber(c) == 1 then" +
            "\nredis.call('expire',KEYS[1],ARGV[2])" +
            "\nend" +
            "\nreturn c;";


    public static HttpServletRequest getHttpServletRequest() {
        return ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();
    }
}
